#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2025/9/8 15:53
# @Author  : Tingkyen
# @File    : 动手搭建transformer模型.py
# @Software: PyCharm
import math

import torch
import torch.nn as nn
from torch.testing._internal.distributed._tensor.common_dtensor import ModelArgs


'''多次头自注意力计算模块'''
class MyMultiHeadAttention(nn.Module):

    def __init__(self,args:ModelArgs,is_causal=False):
        # 构造函数
        #  args：配置对象
        super.__init__()
        assert args.dim % args.heads == 0
        # 模型并⾏处理⼤⼩，默认为1。
        model_parallel_size = 1
        # 本地计算头数，等于总头数除以模型并⾏处理⼤⼩。
        self.n_local_heads = args.n_heads // model_parallel_size
        # 每个头的维度，等于模型维度除以头的总数。
        self.head_dim = args.dim // args.n_heads
        # Wq, Wk, Wv 参数矩阵，每个参数矩阵为 n_embd x n_embd
        # 这⾥通过三个组合矩阵来代替了n个参数矩阵的组合，其逻辑在于矩阵内积再拼接其实等同于拼接矩阵再内积，
        self.wq = nn.Linear(args.dim,self.n_local_heads*self.head_dim,bias=False)
        self.wk = nn.Linear(args.dim,self.n_local_heads*self.head_dim,bias=False)
        self.wv = nn.Linear(args.dim,self.n_local_heads*self.head_dim,bias=False)
        # 输出权重矩阵，维度为 dim x n_embd（head_dim = n_embeds / n_heads）
        self.wo = nn.Linear(self.n_local_heads*self.head_dim,args.dim,bias=False)
        # 注意⼒的 dropout
        self.attention_dropout = nn.Dropout(args.dropout)
        # 残差连接的 dropout
        self.residual_dropout = nn.Dropout(args.dropout)
        # 是否为mask自注意力
        # 创建⼀个上三⻆矩阵，⽤于遮蔽未来信息
        # 注意，因为是多头注意⼒，Mask 矩阵⽐之前我们定义的多⼀个维度
        if is_causal:
            mask = torch.full((1,1,args.max_seq_len,args.max_seq_len),float('-inf'))
            mask = torch.triu(mask,diagonal=1)
            # 注册为模型的缓冲区
            self.register_buffer('mask',mask)

    def forward(self,q:torch.Tensor,k:torch.Tensor,v:torch.Tensor):
        # 获取批次⼤⼩和序列⻓度，[batch_size, seq_len, dim]
        bsz, seqlen, _ = q.shape
        # 计算查询（Q）、键（K）、值（V）,输⼊通过参数矩阵层，维度为 (B, T, n_embed) x (n_embed, n_embed) -> (B, T, n_embed)
        xq, xk, xv = self.wq(q), self.wk(k), self.wv(v)
        # 将 Q、K、V 拆分成多头，维度为 (B, T, n_head, C // n_head)，然后交换维度，变成 (B, n_head, T, C // n_head)
        # 因为在注意⼒计算中我们是取了后两个维度参与计算
        # 为什么要先按B*T*n_head*C//n_head展开再互换1、2维度⽽不是直接按注意⼒输⼊展开，是因为view的展开⽅式是直接把输⼊全部排开，
        # 然后按要求构造，可以发现只有上述操作能够实现我们将每个头对应部分取出来的⽬标
        xq = xq.view(bsz,seqlen,self.n_local_heads,self.head_dim).transpose(1,2)
        xk = xk.view(bsz,seqlen,self.n_local_heads,self.head_dim).transpose(1,2)
        xv = xv.view(bsz,seqlen,self.n_local_heads,self.head_dim).transpose(1,2)

        # 注意⼒计算
        # 计算 QK^T / sqrt(d_k)，维度为 (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
        scores = torch.matmul(xq,xk.transpose(2,3))/math.sqrt(self.head_dim)
        # 掩码⾃注意⼒必须有注意⼒掩码

        # 恢复时间维度并合并头。
        # 将多头的结果拼接起来, 先交换维度为 (B, T, n_head, C // n_head)，再拼接成 (B, T, n_head * C // n_head)
        # contiguous 函数⽤于重新开辟⼀块新内存存储，因为Pytorch设置先transpose再view会报错，
        # 因为view直接基于底层存储得到，然⽽transpose并不会改变底层存储，因此需要额外存储


def print_hi(name):
    print(f"Hi, {name}")


if __name__ == "__main__":
    s = "hello python"
    sl = list(s)
    print(sl)
    print_hi("Tingkyen")
